from tensorflow.keras.applications import ResNet50
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.applications import imagenet_utils
from PIL import Image
import numpy as np
import io
import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "3" 

def prepare_image(image, target):
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    image = image.resize(target)
    image = img_to_array(image)
    image = np.expand_dims(image, axis=0)
    image = imagenet_utils.preprocess_input(image)

    return image 

image = open('./dog.jpg', 'rb').read()
image = Image.open(io.BytesIO(image))
image = prepare_image(image, target=(224, 224))

model = ResNet50(weights='imagenet')
data = {'success': False}
preds = model.predict(image)
results = imagenet_utils.decode_predictions(preds)
data['predictions'] = []

for (imagenetID, label, prob) in results[0]:
    r = {'label': label, 'probability': float(prob)}
    data['predictions'].append(r)

data['success'] = True

print(data)